! Module setting up income
module wage

use params
use normcdf

implicit none

double precision dlw(2), meanshkw(2), norvarw(2), stdevw(2) 

real(prec) newcdfp, newcdfk

double precision, allocatable :: wcontotp(:,:,:), wcontotk(:,:,:),wunden(:), wuntot(:)
double precision, allocatable :: wprobp(:,:,:), wprobk(:,:,:)
real(prec), allocatable :: xw(:,:)
real(prec), allocatable :: incomep(:,:), incomek(:,:)
real(prec), allocatable :: inctrendp(:), inctrendk(:)

double precision, allocatable, private :: meanlw(:,:), sdlw(:,:)

integer, allocatable :: maxwn(:,:), minwn(:,:)

integer, private :: i, t, w, ww
integer, private :: nwtype

double precision, private :: totalp, totalk, phd1,phd2,phd3

contains

! *************************************************
subroutine allocate_wageproc( iret )
    implicit none
    integer, intent(out) :: iret
    integer              :: ier(1)
    
    allocate (incomep(nwp,nt), incomek(nw,nt), &
              inctrendp(nt), inctrendk(nt), &
              wprobp(nwp,nwp,nt), wprobk(nw,nw,nt), &
              maxwn(nw,nt), meanlw(2,nt), minwn(nw,nt), sdlw(2,nt), &
              wcontotp(nwp,nwp,nt),wcontotk(nw,nw,nt), wunden(nw), wuntot(nw), xw(nw,nt),  &
              stat = ier(1))
    
    if ( ier(1) .ne. 0) then
        iret = -1
        return
    endif
    
    iret = 0

end subroutine allocate_wageproc 
! *************************************************
! *************************************************
subroutine wageproc

    !these are taken from ABLNW (pre-estimates)
    double precision, parameter :: age1w = 0.042d0
    double precision, parameter :: age2w = -0.00082d0
    !placeholders for income range
    double precision, parameter :: incmin = 10000d0
    double precision, parameter :: incmax = 100000d0
    double precision :: increment
	
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
    !PARENT INCOME

    !set up income trend
    do t=1,nt
        inctrendp(t) = 1 !constant over time
    enddo

    !determine transition probabilities
    do t=1,nt
        do w=1,nwp
            do ww = 1,nwp
                if (ww.eq.w) then
                    wprobp(ww,w,t) = 1.d0
                    wcontotp(ww,w,t) = 1.d0
                elseif (ww.lt.w) then
                    wprobp(ww,w,t) = 0.d0
                    wcontotp(ww,w,t) = 0.d0
                elseif (ww.gt.w) then
                    wprobp(ww,w,t) = 0.d0
                    wcontotp(ww,w,t) = 1.d0
                endif
            enddo   ! ww
        enddo       ! w
    enddo

    !split income range between nwp nodes
    if (nwp.eq.1) then
        incomep(1,t)=(incmax+incmin)/2.0*inctrendp(t)
    else 
        do t=1,nt
            incomep(1,t)= 8131.1*inctrendp(t)*yrsinpd
            incomep(2,t)=14647.9*inctrendp(t)*yrsinpd
            incomep(3,t)=25006.8*inctrendp(t)*yrsinpd         
        enddo
    endif
    !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

    !KID INCOME

    !1. set up the income trend
    do t=1,nt
        inctrendk(t)=1 !constant income over time
    enddo
    
    if (nw.gt.2) then           
		
		nwtype=nw/2
		
		do i=1,2 !two different variances of income shock
			stdevw(i)   = sqrt(varw(i))
			norvarw(i)	 = varw(i)
			meanshkw(i) = -0.5d0*varw(i)
		enddo
			
        do t=1,nt
        
            !obtain vector for xw (=log w)
            if (rhow.lt.1.d0) then
                meanlw(1,t) = meanshkw(1)*(1.d0-rhow**dble(t))/(1.d0-rhow) 
                sdlw(1,t)   = stdevw(1)*sqrt((1.d0-rhow**(2.d0*dble(t)))/(1.d0-rhow**2.d0))
                meanlw(2,t) = meanshkw(2)*(1.d0-rhow**dble(t))/(1.d0-rhow) 
                sdlw(2,t)   = stdevw(2)*sqrt((1.d0-rhow**(2.d0*dble(t)))/(1.d0-rhow**2.d0))				
            else !unit root!
                meanlw(1,t) = meanshkw(1)*dble(t)
                sdlw(1,t)   = sqrt(dble(t)*varw(1))
                meanlw(2,t) = meanshkw(2)*dble(t)
                sdlw(2,t)   = sqrt(dble(t)*varw(2))				
            endif
            
			xw(1,t)      = meanlw(1,t) + mw*sdlw(1,t)   !biggest log wage
			xw(nwtype,t) = meanlw(1,t) - mw*sdlw(1,t)   !smallest log wage
			dlw(1) = (xw(1,t)-xw(nwtype,t))/dble(nwtype-1)
			do w = nwtype-1,2,-1
				xw(w,t) = xw(w+1,t) + dlw(1)
			enddo
			xw(nwtype+1,t) = meanlw(2,t) + mw*sdlw(2,t)   !biggest log wage
			xw(nw,t) = meanlw(2,t) - mw*sdlw(2,t)   !smallest log wage
			dlw(2) = (xw(nwtype+1,t)-xw(nw,t))/dble(nwtype-1)
			do w = nw-1,nwtype+1,-1
				xw(w,t) = xw(w+1,t) + dlw(2)
			enddo			

            ! find the probabilities for transitions from lw(t-1) to lw(t)
			
			! income type 1
            do w=1,nwtype   !for yesterday
                totalk = 0.d0
                minwn(w,t) = nwtype+1
                maxwn(w,t) = 0
                do ww=1,nwtype  !for today
					
					!first period: will be unconditional
                    if (t.eq.1) then 
                        if (ww.eq.1) then 
                            wprobk(ww,w,t) = 1.d0-normp((xw(ww,t)-0.5d0*dlw(1) - meanshkw(1))/stdevw(1),phd1,phd2,phd3)
                        elseif (ww.eq.nwtype) then
                            wprobk(ww,w,t) = normp((xw(ww,t) + 0.5d0*dlw(1) - meanshkw(1))/stdevw(1),phd1,phd2,phd3)
                        else
                            wprobk(ww,w,t) = normp((xw(ww,t) + 0.5d0*dlw(1) - meanshkw(1))/stdevw(1),phd1,phd2,phd3) - &
                                             normp((xw(ww,t) - 0.5d0*dlw(1) - meanshkw(1))/stdevw(1),phd1,phd2,phd3) 
                        endif
                    
					!later periods: conditional on past income
                    else    !t.gt.1    
                        if (ww.eq.1) then 
                            wprobk(ww,w,t) = 1.d0-normp((xw(ww,t)-0.5d0*dlw(1)-rhow*xw(w,t-1)- &
														 meanshkw(1))/stdevw(1),phd1,phd2,phd3)
                        elseif (ww.eq.nwtype) then
                            wprobk(ww,w,t) = normp((xw(ww,t)+0.5d0*dlw(1)-rhow*xw(w,t-1)- &
												    meanshkw(1))/stdevw(1),phd1,phd2,phd3)
                        else
                            wprobk(ww,w,t) = normp((xw(ww,t)+0.5d0*dlw(1)-rhow*xw(w,t-1)- &
													meanshkw(1))/stdevw(1),phd1,phd2,phd3)- &
                                             normp((xw(ww,t)-0.5d0*dlw(1)-rhow*xw(w,t-1)- &
													meanshkw(1))/stdevw(1),phd1,phd2,phd3)
                        endif
                    endif

                    ! Get min and max wage indices for next period, and adjust probabilities if too small
                    if (wprobk(ww,w,t).lt.1.d-6) then        !if prob transitioning to ww is tiny then:
                        wprobk(ww,w,t) = 0.d0                !...set the probability to zero
                        
                        if (ww.eq.1) then           
                            maxwn(w,t) = ww                 !...set max wage index as that wage's if ww=1 (largest)
                            
                        elseif (ww.gt.1) then
                            if (wprobk(ww-1,w,t).lt.1.d-6) then
                                maxwn(w,t) = ww             !...set max wage index as that wage's if bigger wage also has zero prob
                                                            !...because that means you are on right tail of distn
                            else
                                minwn(w,t) = ww             !...set min wage index as that wage's if bigger wage has larger prob
                                                            !...because that means you are on left tail of distn
                                wprobk(ww+1:nwtype,w,t) = 0.d0   !...and set all lower wages to zero prob
                                exit
                            endif
                        endif           
                    endif
                    
                    totalk = totalk + wprobk(ww,w,t)   !total probability (will be slightly <1 since set tiny probs to zero)
                enddo   !for ww
                
                newcdfk = 0.d0
                do ww=1,nwtype
                    wprobk(ww,w,t) = wprobk(ww,w,t) / totalk   !adjust probabilities to account for total<1
                    newcdfk = newcdfk + wprobk(ww,w,t)
                    wcontotk(ww,w,t) = newcdfk                   
                    newcdfk = newcdfk + 0.d0  
                enddo
                
                if (newcdfk.lt.1.d0-1.0d-7) write(*,*) "total w probability k too low type 1, t= ", t
                if (newcdfk.gt.1.d0+1.0d-7) write(*,*) "total w probability k too high type 1, t= ", t       

            enddo !for w
            
			!unconditional densities and (opposite) CDFs
            if (t.eq.1) then				!so it is just taken from first period
                do w=1,nwtype
           			!unconditional (opposite of) CDF, meaning % of cases above the wage draw
					if (w.eq.nwtype) then		!smallest wage: 100% of cases above it
                        wuntot(w) = 1.d0	
                    else					!bigger wages: subtract off wages smaller than it
                        wuntot(w) = 1.d0 - normp((xw(w,t) - 0.5d0*dlw(1) - meanshkw(1))/stdevw(1),phd1,phd2,phd3)
                    endif
					!unconditional density
                    if (w.eq.1) then
                        wunden(1) = wuntot(1)
                    else
                        wunden(w) = wuntot(w) - wuntot(w-1)
                    endif        
                enddo
            endif
			
			! income type 2
            do w=nwtype+1,nw   !for yesterday
                totalk = 0.d0
                minwn(w,t) = nw+1
                maxwn(w,t) = nwtype
                do ww=nwtype+1,nw  !for today
					
					!first period: will be unconditional
                    if (t.eq.1) then 
                        if (ww.eq.nwtype+1) then 
                            wprobk(ww,w,t) = 1.d0-normp((xw(ww,t)-0.5d0*dlw(2) - meanshkw(2))/stdevw(2),phd1,phd2,phd3)
                        elseif (ww.eq.nw) then
                            wprobk(ww,w,t) = normp((xw(ww,t) + 0.5d0*dlw(2) - meanshkw(2))/stdevw(2),phd1,phd2,phd3)
                        else
                            wprobk(ww,w,t) = normp((xw(ww,t) + 0.5d0*dlw(2) - meanshkw(2))/stdevw(2),phd1,phd2,phd3) - &
                                             normp((xw(ww,t) - 0.5d0*dlw(2) - meanshkw(2))/stdevw(2),phd1,phd2,phd3) 
                        endif
                    
					!later periods: conditional on past income
                    else    !t.gt.1    
                        if (ww.eq.nwtype+1) then 
                            wprobk(ww,w,t) = 1.d0-normp((xw(ww,t)-0.5d0*dlw(2)-rhow*xw(w,t-1)- &
														 meanshkw(2))/stdevw(2),phd1,phd2,phd3)
                        elseif (ww.eq.nw) then
                            wprobk(ww,w,t) = normp((xw(ww,t)+0.5d0*dlw(2)-rhow*xw(w,t-1)- &
													meanshkw(2))/stdevw(2),phd1,phd2,phd3)
                        else
                            wprobk(ww,w,t) = normp((xw(ww,t)+0.5d0*dlw(2)-rhow*xw(w,t-1)- &
													meanshkw(2))/stdevw(2),phd1,phd2,phd3)- &
                                             normp((xw(ww,t)-0.5d0*dlw(2)-rhow*xw(w,t-1)- &
													meanshkw(2))/stdevw(2),phd1,phd2,phd3)
                        endif
                    endif

                    ! Get min and max wage indices for next period, and adjust probabilities if too small
                    if (wprobk(ww,w,t).lt.1.d-6) then        !if prob transitioning to ww is tiny then:
                        wprobk(ww,w,t) = 0.d0                !...set the probability to zero
                        
                        if (ww.eq.nwtype+1) then           
                            maxwn(w,t) = ww                 !...set max wage index as that wage's if ww=1 (largest)
                            
                        elseif (ww.gt.nwtype+1) then
                            if (wprobk(ww-1,w,t).lt.1.d-6) then
                                maxwn(w,t) = ww             !...set max wage index as that wage's if bigger wage also has zero prob
                                                            !...because that means you are on right tail of distn
                            else
                                minwn(w,t) = ww             !...set min wage index as that wage's if bigger wage has larger prob
                                                            !...because that means you are on left tail of distn
                                wprobk(ww+nwtype+1:nw,w,t) = 0.d0   !...and set all lower wages to zero prob
                                exit
                            endif
                        endif           
                    endif
                    
                    totalk = totalk + wprobk(ww,w,t)   !total probability (will be slightly <1 since set tiny probs to zero)
                enddo   !for ww
                
                newcdfk = 0.d0
                do ww=nwtype+1,nw
                    wprobk(ww,w,t) = wprobk(ww,w,t) / totalk   !adjust probabilities to account for total<1
                    newcdfk = newcdfk + wprobk(ww,w,t)
                    wcontotk(ww,w,t) = newcdfk                   
                    newcdfk = newcdfk + 0.d0  
                enddo
                
                if (newcdfk.lt.1.d0-1.0d-7) write(*,*) "total w probability k too low type 2, t= ", t
                if (newcdfk.gt.1.d0+1.0d-7) write(*,*) "total w probability k too high type 2, t= ", t       

            enddo !for w
            
			!unconditional densities and (opposite) CDFs
            if (t.eq.1) then				!so it is just taken from first period
                do w=nwtype+1,nw
           			!unconditional (opposite of) CDF, meaning % of cases above the wage draw
					if (w.eq.nw) then		!smallest wage: 100% of cases above it
                        wuntot(w) = 1.d0	
                    else					!bigger wages: subtract off wages smaller than it
                        wuntot(w) = 1.d0 - normp((xw(w,t) - 0.5d0*dlw(2) - meanshkw(2))/stdevw(2),phd1,phd2,phd3)
                    endif
					!unconditional density
                    if (w.eq.nwtype+1) then
                        wunden(nwtype+1) = wuntot(nwtype+1)
                    else
                        wunden(w) = wuntot(w) - wuntot(w-1)
                    endif        
                enddo
            endif			

        enddo   !over t

    else    !DETERMINISTIC INCOME (nw=2 for 2 income types)
        ! do t=1,nt
            ! !income type 1
            ! wprobk(1,1,t) = 1.d0
            ! wcontotk(1,1,t) = 1.d0
            ! wuntot(1) = 1.d0
            ! xw(1,t) = 0.d0
            ! maxwn(1,t) = nw/2-1
            ! minwn(1,t) = nw/2+1
            ! !income type 2
            ! wprobk(2,2,t) = 1.d0
            ! wcontotk(2,2,t) = 1.d0
            ! wuntot(2) = 1.d0
            ! xw(2,t) = 0.d0
            ! maxwn(2,t) = nw/2
            ! minwn(2,t) = nw+1
        ! enddo
        
    endif   !For nw
        
    !Transform log nodes into actual incomes
    do t=1,nt
           do w=1,nwtype
                incomek(w,t)        = exp(xw(w,t))*inctrendk(t)*meaninc(1) + minw*inctrendk(t)
                incomek(w+nwtype,t) = exp(xw(w+nwtype,t))*inctrendk(t)*meaninc(2) + minw*inctrendk(t)
            enddo
    enddo
	
end subroutine wageproc
! **************************************************

end module wage
